package com.gy.ai.agent.platform.ollama;

import com.alibaba.fastjson2.JSON;
import com.gy.ai.agent.BasicPlatform;
import com.gy.ai.agent.client.NdJsonClient;
import com.gy.ai.agent.client.BaseReqClient;
import com.gy.ai.agent.core.BaseResponse;
import com.gy.ai.agent.core.ReqContext;
import com.gy.ai.agent.core.ReqOptions;
import com.gy.ai.agent.enums.RoleEnums;
import com.gy.ai.agent.listener.BaseReqClientListener;
import com.gy.ai.agent.listener.ReqClientListener;
import com.gy.ai.agent.listener.StreamRespListener;
import com.gy.ai.agent.platform.ollama.entity.OllamaChatReq;
import com.gy.ai.agent.platform.ollama.entity.OllamaChatResp;
import com.gy.ai.agent.platform.ollama.entity.OllamaEmbeddingReq;
import com.gy.ai.agent.platform.ollama.entity.OllamaEmbeddingResp;
import com.gy.ai.agent.util.SysOkHttpClientUtil;

import java.util.Collections;
import java.util.List;

/**
 * @author LGY
 * Create by 2025/6/18 14:31
 */
public class OllamaPlatform extends BasicPlatform<OllamaPlatformConfig> {

    public OllamaPlatform(OllamaPlatformConfig config) {
        super(config);
    }

    // 创建文本对话请求
    public static final String TEXT_CHART = "/api/chat";
    // 创建文本向量请求 （嵌入请求）
    public static final String TEXT_EMBEDDINGS = "/api/embed";

    @Override
    public OllamaChatResp chat(String prompt) {
        OllamaChatReq req = new OllamaChatReq();
        req.setMessages(Collections.singletonList(new OllamaChatReq.Message().setRole(RoleEnums.USER.getValue()).setContent(prompt)));
        return chat(req);
    }

    @Override
    public OllamaChatResp chat(ReqOptions options) {
        OllamaPlatformConfig ollamaPlatformConfig = this.getConfig();
        String url = ollamaPlatformConfig.getEndpoint() + TEXT_CHART;
        // 请求参数
        OllamaChatReq req = (OllamaChatReq) options;
        req.setModel(ollamaPlatformConfig.getModel())
                .setStream(false);
        // 发起请求
        return SysOkHttpClientUtil.doPost(url, ollamaPlatformConfig.getHeaders(), JSON.toJSONString(req), OllamaChatResp.class);
    }

    @Override
    public void chatStream(String prompt, StreamRespListener streamResponseListener) {
        OllamaChatReq req = new OllamaChatReq();
        req.setMessages(Collections.singletonList(new OllamaChatReq.Message().setRole(RoleEnums.USER.getValue()).setContent(prompt)));
        chatStream(req, streamResponseListener);
    }

    @Override
    public void chatStream(ReqOptions options, StreamRespListener streamResponseListener) {
        OllamaPlatformConfig ollamaPlatformConfig = this.getConfig();
        BaseReqClient reqClient = new NdJsonClient();
        String url = ollamaPlatformConfig.getEndpoint() + TEXT_CHART;
        OllamaChatReq req = (OllamaChatReq) options;
        req.setModel(ollamaPlatformConfig.getModel())
                .setStream(true);
        ReqClientListener listener = new BaseReqClientListener(this, streamResponseListener, reqClient, new ReqContext(this, reqClient));
        reqClient.start(url, ollamaPlatformConfig.getHeaders(), JSON.toJSONString(req), listener);
    }

    @Override
    public OllamaEmbeddingResp embedding(String input) {
        OllamaEmbeddingReq req = new OllamaEmbeddingReq();
        req.setInput(input);
        return embedding(req);
    }

    @Override
    public OllamaEmbeddingResp embedding(ReqOptions options) {
        OllamaPlatformConfig ollamaPlatformConfig = this.getConfig();
        String url = ollamaPlatformConfig.getEndpoint() + TEXT_EMBEDDINGS;
        // 请求参数
        OllamaEmbeddingReq req = (OllamaEmbeddingReq) options;
        req.setModel(ollamaPlatformConfig.getModel());
        // 发起请求
        return SysOkHttpClientUtil.doPost(url, ollamaPlatformConfig.getHeaders(), JSON.toJSONString(req), OllamaEmbeddingResp.class);
    }

    @Override
    public BaseResponse rerank(String query, List<String> documents) {
        return rerank(ReqOptions.DEFAULT);
    }

    @Override
    public BaseResponse rerank(ReqOptions options) {
        throw new RuntimeException("当前平台暂不支持重排");
    }

    @Override
    public BaseResponse generateImage(String prompt) {
        return generateImage(ReqOptions.DEFAULT);
    }

    @Override
    public BaseResponse generateImage(ReqOptions options) {
        throw new RuntimeException("当前平台暂不支持文生图");
    }
}
